package com.yale.elasticsearch.service.aggregation;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.aggregations.*;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.json.JsonData;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.List;
import java.util.Map;

/**
 * 聚合 Bucket
 */

@Service
@Slf4j
public class AggrBucketService {

    @Resource
    private ElasticsearchClient elasticsearchClient;

    /**
     * 按岁数进行聚合分桶
     */
    public Object aggrBucketTerms() {
        String responseResult = "";
        try {
            // 创建Terms聚合，按age字段进行分桶
            Aggregation ageTermsAgg = AggregationBuilders.terms()
                    .field("age")
                    .size(10)
                    .build()._toAggregation();

            // 创建聚合查询请求
            SearchRequest searchRequest = SearchRequest.of(s -> s
                    .index("mydlq-user")
                    .size(0) // 不返回文档，只返回聚合结果
                    .aggregations("age_bucket", ageTermsAgg)
            );

            // 执行请求
            SearchResponse<JsonData> searchResponse = elasticsearchClient.search(searchRequest, JsonData.class);

            // 获取聚合结果
            Map<String, Aggregate> aggregations = searchResponse.aggregations();
            Aggregate ageBucketAgg = aggregations.get("age_bucket");

            if (ageBucketAgg != null && ageBucketAgg.isSterms()) {
                StringTermsAggregate termsAggregate = ageBucketAgg.sterms();
                termsAggregate.buckets().array().forEach(bucket -> {
                    System.out.println("桶名：" + bucket.key() + ",  | 总数：" +bucket.docCount());
                });
            }

            // 返回响应对象Json串
            responseResult = searchResponse.toString();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return responseResult;

    }

    /**
     * 按工资范围进行聚合分桶
     */
    public Object aggrBucketRange() {
        String responseResult = "";
        try {
            // 创建 Range 聚合，按 salary 字段进行分桶
            Aggregation rangeAgg = new RangeAggregation.Builder()
                    .field("salary")
                    .ranges(
                            new AggregationRange.Builder().from(0.0).to(3000.0).key("低级员工").build(),
                            new AggregationRange.Builder().from(3000.0).to(5000.0).key("中级员工").build(),
                            new AggregationRange.Builder().from(5000.0).to(9000.0).key("高级员工").build()
                    )
                    .build()._toAggregation();

            // 创建聚合查询请求
            SearchRequest searchRequest = SearchRequest.of(s -> s
                    .index("mydlq-user")
                    .size(0) // 不返回文档，只返回聚合结果
                    .aggregations("salary_range_bucket", rangeAgg)
            );

            // 执行请求
            SearchResponse<JsonData> searchResponse = elasticsearchClient.search(searchRequest, JsonData.class);

            // 获取聚合结果
            RangeAggregate rangeAggregate = searchResponse.aggregations().get("salary_range_bucket").range();

            // 输出内容
            if (rangeAggregate != null) {
                for (RangeBucket bucket : rangeAggregate.buckets().array()) {
                    System.out.println("桶名：" + bucket.key() + " | 总数：" + bucket.docCount());
                }
            }

            // 返回响应对象Json串
            responseResult = searchResponse.toString();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return responseResult;

    }

    /**
     * 按照时间范围进行分桶
     */
    public Object aggrBucketDateRange() {
        String responseResult = "";
        try {
            // 创建 DateRange 聚合
            Aggregation dateRangeAgg = new DateRangeAggregation.Builder()
                    .field("birthDate")
                    .format("yyyy-MM-dd")
                    .ranges( new DateRangeExpression.Builder()
                                    .key("出生日期1985-1990的员工")
                                    .from(FieldDateMath.of(builder -> builder.expr("1985-01-01")))
                                    .to(FieldDateMath.of(builder -> builder.expr("1990-12-31")))
                                    .build(),
                            new DateRangeExpression.Builder()
                                    .key("出生日期1990-1995的员工")
                                    .from(FieldDateMath.of(builder -> builder.expr("1990-01-01")))
                                    .to(FieldDateMath.of(builder -> builder.expr("1995-12-31")))
                                    .build())
                    .build()._toAggregation();

            // 创建聚合查询请求
            SearchRequest searchRequest = SearchRequest.of(s -> s
                    .index("mydlq-user")
                    .size(0) // 不返回文档，只返回聚合结果
                    .aggregations("date_range_bucket", dateRangeAgg)
            );

            // 执行请求
            SearchResponse<JsonData> searchResponse = elasticsearchClient.search(searchRequest, JsonData.class);

            // 获取聚合结果
            Aggregate dateRangeBucket = searchResponse.aggregations().get("date_range_bucket");

            // 输出内容
            if (dateRangeBucket != null) {
                dateRangeBucket.dateRange().buckets().array().forEach(bucket -> {
                    System.out.println("桶名：" + bucket.key() + " | 总数：" + bucket.docCount());
                });
            }

            // 返回响应对象Json串
            responseResult = searchResponse.toString();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return responseResult;

    }

    /**
     * 按工资多少进行聚合分桶
     */
    public Object aggrBucketHistogram() {
        String responseResult = "";
        try {
            // 构建聚合查询
            SearchRequest searchRequest = SearchRequest.of(s -> s
                    .index("mydlq-user")
                    .size(0)
                    .aggregations("salary_histogram", a -> a.histogram(h -> h
                            .field("salary")
                            .interval(3000.0)
                            .extendedBounds(eb -> eb.min(0.0).max(12000.0))
                    ))
            );

            // 执行请求
            SearchResponse<Void> response = elasticsearchClient.search(searchRequest, Void.class);

            // 获取聚合结果
            Aggregate aggregate = response.aggregations().get("salary_histogram");

            if (aggregate != null && aggregate.isHistogram()) {
                HistogramAggregate salaryHistogram = aggregate.histogram();

                // 输出内容
                List<HistogramBucket> buckets = salaryHistogram.buckets().array();
                System.out.println("-------------------------------------------");
                System.out.println("聚合信息:");
                for (HistogramBucket bucket : buckets) {
                    System.out.printf("桶名：%s | 总数：%d%n", bucket.key(), bucket.docCount());
                }
                System.out.println("-------------------------------------------");

                responseResult = response.toString();
            }
        } catch (IOException e) {
            System.err.println("Error executing search request: " + e.getMessage());
        }
        return responseResult;

    }

    /**
     * 按出生日期进行分桶
     */
    public Object aggrBucketDateHistogram() {
        String responseResult = "";
        try {
            // 创建 date_histogram 聚合
            DateHistogramAggregation dateHistogramAgg = DateHistogramAggregation.of(dh -> dh
                    .field("birthDate")
                    .calendarInterval(CalendarInterval.Year) // 按年分桶
                    .format("yyyy")
            );

            // 创建聚合查询请求
            SearchRequest searchRequest = SearchRequest.of(s -> s
                    .index("mydlq-user")
                    .size(0) // 不返回文档，只返回聚合结果
                    .aggregations("birthday_histogram", a -> a.dateHistogram(dateHistogramAgg))
            );

            // 执行请求
            SearchResponse<Void> response = elasticsearchClient.search(searchRequest, Void.class);

            // 获取聚合结果
            Aggregate aggregate = response.aggregations().get("birthday_histogram");

            if (aggregate != null && aggregate.isDateHistogram()) {
                DateHistogramAggregate dateHistogram = aggregate.dateHistogram();

                // 输出内容
                List<DateHistogramBucket> buckets = dateHistogram.buckets().array();
                System.out.println("-------------------------------------------");
                System.out.println("聚合信息:");
                for (DateHistogramBucket bucket : buckets) {
                    System.out.printf("桶名：%s | 总数：%d%n", bucket.keyAsString(), bucket.docCount());
                }
                System.out.println("-------------------------------------------");

                responseResult = response.toString();
            }
        } catch (IOException e) {
            System.err.println("Error executing search request: " + e.getMessage());
        }
        return responseResult;
    }


}
